from rest_framework.generics import ListAPIView
from rest_framework.viewsets import ViewSet
from demo.utils import pre_precipitaion, pre_temperature
from .serializers import TrainModelSerializer, PredictModelSerializer
from .models import Train, Predict
from rest_framework.response import Response
from .utils import utils, axis, model_load
from rest_framework import status


# Create your views here.


class TrainListAPIView(ListAPIView):
    """
    训练视图
    """
    queryset = Train.objects.filter()
    serializer_class = TrainModelSerializer


class PredictListAPIView(ListAPIView):
    """
    预测数据视图
    """
    queryset = Predict.objects.filter()
    serializer_class = PredictModelSerializer


class PredictAPIView(ViewSet):
    """
    预测视图
    """

    def pre_pre(self, request):
        flag = 1
        model = model_load('demo/libs/model-降水.h5')
        time = request.data.get('time')
        pre_pre = pre_precipitaion.to_pre(time, model)
        date = axis()
        result = utils(time, flag)

        return Response({
            'msg': '预测成功',
            'data': pre_pre,
            'data_': result,
            'date': date,
        }, status=status.HTTP_200_OK)

    def pre_tem(self, request):
        flag = 2
        model = model_load('demo/libs/model-气温.h5')
        time = request.data.get('time')
        print(time)
        pre_tem = pre_temperature.to_pre(time, model)
        date = axis()
        result = utils(time, flag)
        print(pre_tem)

        return Response({
            'msg': '预测成功',
            'data': pre_tem,
            'data_': result,
            'date': date,
        }, status=status.HTTP_200_OK)
